package com.doubao.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.service.ArkService;
import com.volcengine.ark.runtime.exception.ArkHttpException;
import lombok.Data;
import javax.annotation.PostConstruct;
import java.util.*;
import java.time.Duration;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("/api/chat")
@CrossOrigin(origins = "*")
public class ChatController {

    private ArkService arkService;

    // 使用 ConcurrentHashMap 存储每个会话的历史记录
    private final Map<String, List<ChatMessage>> chatHistory = new ConcurrentHashMap<>();

    @Value("${doubao.api.endpoint-id}")
    private String endpointId;

    @Value("${doubao.api.api-key}")
    private String apiKey;

    // 最大记忆轮数
    private static final int MAX_HISTORY_TURNS = 10;

    @PostConstruct
    public void init() {
        log.info("正在初始化ChatController, endpointId: {}", endpointId);

        this.arkService = ArkService.builder()
                .apiKey(apiKey)
                .baseUrl("https://ark.cn-beijing.volces.com/api/v3")
                .timeout(Duration.ofSeconds(60))
                .retryTimes(3)
                .build();
    }

    @Data
    public static class ChatRequest {
        private String message;
        private String sessionId;  // 添加会话ID
    }

    @Data
    public static class ChatResponse {
        private String status;
        private String message;
        private String sessionId;  // 返回会话ID

        public ChatResponse(String status, String message, String sessionId) {
            this.status = status;
            this.message = message;
            this.sessionId = sessionId;
        }
    }

    private String generateSessionId() {
        return UUID.randomUUID().toString();
    }

    private List<ChatMessage> getOrCreateHistory(String sessionId) {
        return chatHistory.computeIfAbsent(sessionId, k -> new ArrayList<>());
    }

    @PostMapping("/send")
    public ResponseEntity<ChatResponse> sendMessage(@RequestBody ChatRequest request) {
        try {
            // 获取或生成会话ID
            String sessionId = (request.getSessionId() != null) ?
                              request.getSessionId() :
                              generateSessionId();

            log.info("收到聊天请求: sessionId={}, message={}", sessionId, request.getMessage());

            // 获取会话历史
            List<ChatMessage> history = getOrCreateHistory(sessionId);

            // 构建完整的消息列表
            List<ChatMessage> messages = new ArrayList<>();

            // 添加系统消息
            messages.add(ChatMessage.builder()
                    .role(ChatMessageRole.SYSTEM)
                    .content("你是豆包，是由字节跳动开发的 AI 人工智能助手。请记住用户之前说过的话，保持对话连贯。")
                    .build());

            // 添加历史消息
            messages.addAll(history);

            // 添加新的用户消息
            ChatMessage userMessage = ChatMessage.builder()
                    .role(ChatMessageRole.USER)
                    .content(request.getMessage())
                    .build();
            messages.add(userMessage);

            // 创建请求
            ChatCompletionRequest chatRequest = ChatCompletionRequest.builder()
                    .model(endpointId)
                    .messages(messages)
                    .build();

            // 发送请求并获取响应
            ChatCompletionResult result = arkService.createChatCompletion(chatRequest);

            if (result != null && result.getChoices() != null && !result.getChoices().isEmpty()) {
                ChatCompletionChoice choice = result.getChoices().get(0);
                ChatMessage assistantMessage = choice.getMessage();
                String reply = (String) assistantMessage.getContent();

                // 更新历史记录
                history.add(userMessage);
                history.add(assistantMessage);

                // 如果历史记录太长，删除较早的对话
                while (history.size() > MAX_HISTORY_TURNS * 2) {
                    history.remove(0);
                    history.remove(0);
                }

                log.info("收到AI回复: sessionId={}, reply={}", sessionId, reply);
                return ResponseEntity.ok(new ChatResponse("success", reply, sessionId));
            }

            return ResponseEntity.ok(new ChatResponse("error", "AI未返回有效回复", sessionId));

        } catch (ArkHttpException e) {
            log.error("调用豆包AI服务失败: {}", e.getMessage(), e);
            return ResponseEntity.ok(new ChatResponse(
                "error",
                "服务调用失败: " + e.getMessage(),
                request.getSessionId()
            ));
        } catch (Exception e) {
            log.error("处理聊天请求时发生错误", e);
            return ResponseEntity.ok(new ChatResponse(
                "error",
                "系统内部错误: " + e.getMessage(),
                request.getSessionId()
            ));
        }
    }

    // 清除会话历史
    @DeleteMapping("/session/{sessionId}")
    public ResponseEntity<ChatResponse> clearSession(@PathVariable String sessionId) {
        chatHistory.remove(sessionId);
        return ResponseEntity.ok(new ChatResponse(
            "success",
            "会话历史已清除",
            sessionId
        ));
    }

    // 获取会话历史
    @GetMapping("/session/{sessionId}/history")
    public ResponseEntity<List<String>> getSessionHistory(@PathVariable String sessionId) {
        List<ChatMessage> history = chatHistory.getOrDefault(sessionId, new ArrayList<>());
        List<String> messages = history.stream()
        .map(msg -> msg.getRole() + ": " + msg.getContent())
        .collect(Collectors.toList());
        return ResponseEntity.ok(messages);
    }
}